#coding = utf-8

'''
将数据转换成chengkung
'''

import os
from pathlib2 import Path
import SimpleITK as sitk
import numpy as np

def normalize(vol):
    hu_max = 250
    hu_min = -200
    vol = np.clip(vol, hu_min, hu_max)

    mxval = np.max(vol)
    mnval = np.min(vol)
    volume_norm = (vol - mnval) / max(mxval - mnval, 1e-3)

    return volume_norm

def normalizeVessle(vol):
    hu_max = 600
    hu_min = 60
    vol = np.clip(vol, hu_min, hu_max)

    mxval = np.max(vol)
    mnval = np.min(vol)
    volume_norm = (vol - mnval) / max(mxval - mnval, 1e-3)

    return volume_norm

def read_data_and_mask():
    data_path = "/datasets/qiye/DongBeiDaXue2/image_arterial"
    mask_path = "/datasets/qiye/DongBeiDaXue2/artery"
    save_path = "/datasets/DongbeiDaxue/chengkung_artery"

    data_file_list = sorted(os.listdir(data_path))
    mask_file_list = sorted(os.listdir(mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        mask_file_name = mask_file_list[i]
        assert data_file_name.split("_")[1] == mask_file_name.split("_")[1], data_file_name+","+mask_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        mask_file_name = os.path.join(mask_path, mask_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i+50).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        mask = sitk.GetArrayFromImage(sitk.ReadImage(mask_file_name))

        assert data.shape == mask.shape, str(data.shape[0]) + "," + str(mask.shape[0]) + "," + data_file_name
        assert len(np.unique(mask)) <= 2, np.unique(mask)
        print(np.unique(mask))



        for j in range(data.shape[0]):
            data_item = normalizeVessle(data[j]).astype(np.float32)
            mask_item = mask[j]
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), data_item)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), mask_item)

        print("finish {}".format(i))


def read_data_and_mask_tumor():
    data_path = "/datasets/qiye/DongBeiDaXue2/image_venous"
    liver_mask_path = "/datasets/qiye/DongBeiDaXue2/liver"
    tumor_mask_path = "/datasets/qiye/DongBeiDaXue2/lesion"
    save_path = "/datasets/DongbeiDaxue/chengkunv2"

    data_file_list = sorted(os.listdir(data_path))
    liver_file_list = sorted(os.listdir(liver_mask_path))
    tumor_file_list = sorted(os.listdir(tumor_mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        liver_file_name = liver_file_list[i]
        tumor_file_name = tumor_file_list[i]
        assert data_file_name.split("_")[1] == liver_file_name.split("_")[1], data_file_name+","+liver_file_name
        assert data_file_name.split("_")[1] == tumor_file_name.split("_")[1], data_file_name + "," + tumor_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        liver_file_name = os.path.join(liver_mask_path, liver_file_name)
        tumor_file_name = os.path.join(tumor_mask_path, tumor_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i+50).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        liver = sitk.GetArrayFromImage(sitk.ReadImage(liver_file_name))
        tumor = sitk.GetArrayFromImage(sitk.ReadImage(tumor_file_name))

        assert data.shape == liver.shape, str(data.shape[0]) + "," + str(liver.shape[0]) + "," + data_file_name
        assert data.shape == tumor.shape, str(data.shape[0]) + "," + str(tumor.shape[0]) + "," + data_file_name
        assert len(np.unique(liver)) <= 2, np.unique(liver)
        print(np.unique(liver), liver.sum())
        assert len(np.unique(tumor)) <= 2, np.unique(tumor)
        print(np.unique(tumor), tumor.sum())



        for j in range(data.shape[0]):
            data_item = normalize(data[j]).astype(np.float32)
            mask_item = liver[j]
            mask_item[tumor[j] == 1] = 2
            #print(np.unique(mask_item))
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), data_item)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), mask_item)


        print("finish {}".format(i))
        print()

#只保留肝脏的切片
def read_data_and_mask_tumor_v2():
    data_path = "/datasets/qiye/DongBeiDaXue2/image_venous"
    liver_mask_path = "/datasets/qiye/DongBeiDaXue2/liver"
    tumor_mask_path = "/datasets/qiye/DongBeiDaXue2/lesion"
    save_path = "/datasets/DongbeiDaxue/chengkun_only_liver"

    data_file_list = sorted(os.listdir(data_path))
    liver_file_list = sorted(os.listdir(liver_mask_path))
    tumor_file_list = sorted(os.listdir(tumor_mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        liver_file_name = liver_file_list[i]
        tumor_file_name = tumor_file_list[i]
        assert data_file_name.split("_")[1] == liver_file_name.split("_")[1], data_file_name+","+liver_file_name
        assert data_file_name.split("_")[1] == tumor_file_name.split("_")[1], data_file_name + "," + tumor_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        liver_file_name = os.path.join(liver_mask_path, liver_file_name)
        tumor_file_name = os.path.join(tumor_mask_path, tumor_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i+50).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        liver = sitk.GetArrayFromImage(sitk.ReadImage(liver_file_name))
        tumor = sitk.GetArrayFromImage(sitk.ReadImage(tumor_file_name))

        assert data.shape == liver.shape, str(data.shape[0]) + "," + str(liver.shape[0]) + "," + data_file_name
        assert data.shape == tumor.shape, str(data.shape[0]) + "," + str(tumor.shape[0]) + "," + data_file_name
        assert len(np.unique(liver)) <= 2, np.unique(liver)
        print(np.unique(liver), liver.sum())
        assert len(np.unique(tumor)) <= 2, np.unique(tumor)
        print(np.unique(tumor), tumor.sum())


        count = 0
        for j in range(data.shape[0]):
            data_item = normalize(data[j]).astype(np.float32)
            mask_item = liver[j]
            mask_item[tumor[j] == 1] = 2
            if np.max(mask_item) == 0:
                continue
            #print(np.unique(mask_item))
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(count).zfill(3))), mask_item)
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(count).zfill(3))),
                    data_item)

            #mask_item[mask_item == 2] = 1
            #np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(count).zfill(3))), data_item*mask_item)
            count += 1


        print("finish {}".format(i))
        print()

#只保留肿瘤的切片
def read_data_and_mask_tumor_v3():
    data_path = "/datasets/qiye/DongBeiDaXue2/image_venous"
    liver_mask_path = "/datasets/qiye/DongBeiDaXue2/liver"
    tumor_mask_path = "/datasets/qiye/DongBeiDaXue2/lesion"
    save_path = "/datasets/DongbeiDaxue/chengkun_only_tumor"

    data_file_list = sorted(os.listdir(data_path))
    liver_file_list = sorted(os.listdir(liver_mask_path))
    tumor_file_list = sorted(os.listdir(tumor_mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        liver_file_name = liver_file_list[i]
        tumor_file_name = tumor_file_list[i]
        assert data_file_name.split("_")[1] == liver_file_name.split("_")[1], data_file_name+","+liver_file_name
        assert data_file_name.split("_")[1] == tumor_file_name.split("_")[1], data_file_name + "," + tumor_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        liver_file_name = os.path.join(liver_mask_path, liver_file_name)
        tumor_file_name = os.path.join(tumor_mask_path, tumor_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i+50).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        liver = sitk.GetArrayFromImage(sitk.ReadImage(liver_file_name))
        tumor = sitk.GetArrayFromImage(sitk.ReadImage(tumor_file_name))

        assert data.shape == liver.shape, str(data.shape[0]) + "," + str(liver.shape[0]) + "," + data_file_name
        assert data.shape == tumor.shape, str(data.shape[0]) + "," + str(tumor.shape[0]) + "," + data_file_name
        assert len(np.unique(liver)) <= 2, np.unique(liver)
        print(np.unique(liver), liver.sum())
        assert len(np.unique(tumor)) <= 2, np.unique(tumor)
        print(np.unique(tumor), tumor.sum())


        count = 0
        for j in range(data.shape[0]):
            data_item = normalize(data[j]).astype(np.float32)
            mask_item = liver[j]
            mask_item[tumor[j] == 1] = 2
            if np.max(mask_item) <= 1:
                continue
            #print(np.unique(mask_item))
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(count).zfill(3))), mask_item)
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(count).zfill(3))),
                    data_item)

            #mask_item[mask_item == 2] = 1
            #np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(count).zfill(3))), data_item*mask_item)
            count += 1


        print("finish {}".format(i))
        print()

def read_data_and_mask_tumor3D():
    data_path = "/datasets/qiye/image_venous"
    liver_mask_path = "/datasets/qiye/liver"
    tumor_mask_path = "/datasets/qiye/lesion"
    save_path = "/datasets/DongbeiDaxue/chengkun3D"

    data_file_list = sorted(os.listdir(data_path))
    liver_file_list = sorted(os.listdir(liver_mask_path))
    tumor_file_list = sorted(os.listdir(tumor_mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        liver_file_name = liver_file_list[i]
        tumor_file_name = tumor_file_list[i]
        assert data_file_name.split("_")[1] == liver_file_name.split("_")[1], data_file_name+","+liver_file_name
        assert data_file_name.split("_")[1] == tumor_file_name.split("_")[1], data_file_name + "," + tumor_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        liver_file_name = os.path.join(liver_mask_path, liver_file_name)
        tumor_file_name = os.path.join(tumor_mask_path, tumor_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        liver = sitk.GetArrayFromImage(sitk.ReadImage(liver_file_name))
        tumor = sitk.GetArrayFromImage(sitk.ReadImage(tumor_file_name))

        assert data.shape == liver.shape, str(data.shape[0]) + "," + str(liver.shape[0]) + "," + data_file_name
        assert data.shape == tumor.shape, str(data.shape[0]) + "," + str(tumor.shape[0]) + "," + data_file_name
        assert len(np.unique(liver)) <= 2, np.unique(liver)
        print(np.unique(liver), liver.sum())
        assert len(np.unique(tumor)) <= 2, np.unique(tumor)
        print(np.unique(tumor), tumor.sum())


        data_list = []
        mask_list = []
        index = 0
        for j in range(data.shape[0]):
            data_item = normalize(data[j]).astype(np.float32)
            mask_item = liver[j]
            mask_item[tumor[j] == 1] = 2
            data_list.append(data_item)
            mask_list.append(mask_item)
            if j >0 and (j+1) % 16 == 0:
                data_np = np.array(data_list)
                mask_np = np.array(mask_list)
                data_list.clear()
                mask_list.clear()
                index = int((j + 1) / 16) - 1
                data_np1 = data_np[:, 0:256, 0:256]
                mask_np1 = mask_np[:, 0:256, 0:256]
                np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index*4).zfill(3))), data_np1)
                np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index*4).zfill(3))), mask_np1)
                data_np2 = data_np[:, 256:, 0:256]
                mask_np2 = mask_np[:, 256:, 0:256]
                np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 1).zfill(3))), data_np2)
                np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 1).zfill(3))), mask_np2)
                data_np3 = data_np[:, 0:256, 256:]
                mask_np3 = mask_np[:, 0:256, 256:]
                np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 2).zfill(3))), data_np3)
                np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 2).zfill(3))), mask_np3)
                data_np4 = data_np[:, 256:, 256:]
                mask_np4 = mask_np[:, 256:, 256:]
                np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 3).zfill(3))), data_np4)
                np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 3).zfill(3))), mask_np4)




        if len(data_list) > 0:
            for s in range(len(data_list), 16):
                data_list.append(data_list[-1])
                mask_list.append(mask_list[-1])
            data_np = np.array(data_list)
            mask_np = np.array(mask_list)
            index += 1
            data_np1 = data_np[:, 0:256, 0:256]
            mask_np1 = mask_np[:, 0:256, 0:256]
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4).zfill(3))), data_np1)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4).zfill(3))), mask_np1)
            data_np2 = data_np[:, 256:, 0:256]
            mask_np2 = mask_np[:, 256:, 0:256]
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 1).zfill(3))), data_np2)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 1).zfill(3))), mask_np2)
            data_np3 = data_np[:, 0:256, 256:]
            mask_np3 = mask_np[:, 0:256, 256:]
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 2).zfill(3))), data_np3)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 2).zfill(3))), mask_np3)
            data_np4 = data_np[:, 256:, 256:]
            mask_np4 = mask_np[:, 256:, 256:]
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index * 4 + 3).zfill(3))), data_np4)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index * 4 + 3).zfill(3))), mask_np4)


            #print(data_item.shape, mask_item.shape)
            #print(np.unique(mask_item))
            #np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), data_item)
            #np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), mask_item)


        print("finish {}".format(i))
        print()

def read_data_and_mask_tumor_venous():
    data_path = "/datasets/qiye/DongBeiDaXue2/image_venous"
    liver_mask_path = "/datasets/qiye/DongBeiDaXue2/liver"
    tumor_mask_path = "/datasets/qiye/DongBeiDaXue2/lesion"
    hepatic_mask_path = "/datasets/qiye/DongBeiDaXue2/hepatic_vein"
    portal_mask_path = "/datasets/qiye/DongBeiDaXue2/portal_vein"
    save_path = "/datasets/DongbeiDaxue/chengkun_join"

    data_file_list = sorted(os.listdir(data_path))
    liver_file_list = sorted(os.listdir(liver_mask_path))
    tumor_file_list = sorted(os.listdir(tumor_mask_path))
    hepatic_file_list = sorted(os.listdir(hepatic_mask_path))
    portal_file_list = sorted(os.listdir(portal_mask_path))
    for i in range(len(data_file_list)):
        data_file_name = data_file_list[i]
        liver_file_name = liver_file_list[i]
        tumor_file_name = tumor_file_list[i]
        hepatic_file_name = hepatic_file_list[i]
        portal_file_name = portal_file_list[i]
        assert data_file_name.split("_")[1] == liver_file_name.split("_")[1], data_file_name + "," + liver_file_name
        assert data_file_name.split("_")[1] == tumor_file_name.split("_")[1], data_file_name + "," + tumor_file_name
        assert data_file_name.split("_")[1] == hepatic_file_name.split("_")[1], data_file_name + "," + hepatic_file_name
        assert data_file_name.split("_")[1] == portal_file_name.split("_")[1], data_file_name + "," + portal_file_name
        data_file_name = os.path.join(data_path, data_file_name)
        liver_file_name = os.path.join(liver_mask_path, liver_file_name)
        tumor_file_name = os.path.join(tumor_mask_path, tumor_file_name)
        hepatic_file_name = os.path.join(hepatic_mask_path, hepatic_file_name)
        portal_file_name = os.path.join(portal_mask_path, portal_file_name)
        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i+50).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)




        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        data = sitk.GetArrayFromImage(sitk.ReadImage(data_file_name))
        liver = sitk.GetArrayFromImage(sitk.ReadImage(liver_file_name))
        tumor = sitk.GetArrayFromImage(sitk.ReadImage(tumor_file_name))
        hepatic_vein = sitk.GetArrayFromImage(sitk.ReadImage(hepatic_file_name))
        portal_vein = sitk.GetArrayFromImage(sitk.ReadImage(portal_file_name))

        assert data.shape == liver.shape, str(data.shape[0]) + "," + str(liver.shape[0]) + "," + data_file_name
        assert data.shape == tumor.shape, str(data.shape[0]) + "," + str(tumor.shape[0]) + "," + data_file_name
        assert data.shape == hepatic_vein.shape, str(data.shape[0]) + "," + str(hepatic_vein.shape[0]) + "," + data_file_name
        assert data.shape == portal_vein.shape, str(data.shape[0]) + "," + str(portal_vein.shape[0]) + "," + data_file_name
        assert len(np.unique(liver)) <= 2, np.unique(liver)
        print(np.unique(liver), liver.sum())
        assert len(np.unique(tumor)) <= 2, np.unique(tumor)
        print(np.unique(tumor), tumor.sum())
        assert len(np.unique(hepatic_vein)) <= 2, np.unique(hepatic_vein)
        print(np.unique(hepatic_vein), hepatic_vein.sum())
        assert len(np.unique(portal_vein)) <= 2, np.unique(portal_vein)
        print(np.unique(portal_vein), portal_vein.sum())

        for j in range(data.shape[0]):
            data_item = normalize(data[j]).astype(np.float32)
            mask_item = liver[j]
            mask_item[tumor[j] == 1] = 2
            mask_item[hepatic_vein[j] == 1] = 3
            mask_item[portal_vein[j] == 1] = 4
            # print(np.unique(mask_item))
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), data_item)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), mask_item)

        print("finish {}".format(i))
        print()



if __name__ == '__main__':
    read_data_and_mask_tumor()

